"""

Computational Economics
2.3: Basic Plotting with Matplotlib
http://johnstachurski.net/lectures/basicplot.html

    This file contains daily quotes for the Nikkei 225 from Jan 1984 until
    May 2009, downloaded from Yahoo finance

    Here are the first few lines

    Date,Open,High,Low,Close,Volume,Adj Close
    2009-05-21,9280.35,9286.35,9189.92,9264.15,133200,9264.15
    2009-05-20,9372.72,9399.40,9311.61,9344.64,143200,9344.64
    2009-05-19,9172.56,9326.75,9166.97,9290.29,167000,9290.29
    2009-05-18,9167.05,9167.82,8997.74,9038.69,147800,9038.69
    2009-05-15,9150.21,9272.08,9140.90,9265.02,172000,9265.02

    Data is comma separated (csv), with most recent date first

    For our price data we will use the last column (Adj Close)

"""

from os.path import (dirname, join)
import pylab
from datetime import date

THIS_DIR    = dirname(__file__)
FNAME       = 'nikkei.csv'
DATA_PATH   = join(THIS_DIR, FNAME)


# Exercise 1
def ex1():
    """
    Exercise 1:

    Plot the data (i.e., the Adj Close column) as a time series

        * Use the File I/O operations in this lecture to extract the data
              o You might like to use the string method split()
              o Note that there is a module called csv for working with csv
                files
                + But don't use it this time: I want you to practice basic file
                  I/O
        * Make sure your time series is from earliest (i.e., Jan 84) to latest
          (i.e., May 2009)
    """
    # load data
    f = open(DATA_PATH)
    lines = f.readlines()
    f.close()

    # organize data
    dates, closings = [], []
    d = 0
    for l in lines[1:]:
        field = l.split(',')
        d += 1
        dates.append(d)
        closings.append(field[-1])

    # plot
    #print zip(dates, closings)
    pylab.plot(dates, closings)
    pylab.show()


# Exercise 2
daily_return = lambda t, y: ((float(t) - float(y)) / float(y)) * 100


def ex2():
    """
    Exercise 2:

    Write a function that

        * takes a start year and an end year, and
        * plots daily returns (as a percentage)

    Daily return = [(today - yesterday) / yesterday] * 100
    """
    # get start and end years
    start_year = int(raw_input("Enter start year (between 1984 and 2009): "))
    if start_year < 1984:
        start_year = 1984
    if start_year > 2009:
        start_year = 2009

    end_year = int(raw_input("Enter end year: "))
    if end_year < start_year:
        start_year, end_year = end_year, start_year


    # load data
    f = open(DATA_PATH)
    lines = f.readlines()
    f.close()

    # organize data
    dates, returns = [], []
    for n in range(2, len(lines)):
        # parse data
        tl = lines[n]
        yl = lines[n-1]
        today = tl.split(',')
        year, m, d = today[0].split('-')

        # check range
        if int(year) < start_year or int(year) > end_year:
            continue

        yesterday = yl.split(',')
        r = daily_return(today[-1], yesterday[-1])

        dates.append(n)
        returns.append(r)

    # plot
    #print zip(dates, closings)
    pylab.plot(dates, returns)
    pylab.xlabel('years %s to %s' % (start_year, end_year))
    pylab.ylabel('daily % return')
    pylab.title('Nikkei Index Daily Returns')
    pylab.show()


# Exercise 3
def read_data_file(path):
    """discard first line as header"""
    f = open(DATA_PATH)
    lines = f.readlines()
    f.close()
    return lines[1:]


def get_daily_returns(csv_lines):
    """data should be a list of lines from csv file with date first item, close
    as last. Return a list of tuples with items: (date, % change)"""
    returns = []

    for n in range(1, len(csv_lines)):
        # parse data
        tl = csv_lines[n]
        yl = csv_lines[n-1]
        today = tl.split(',')
        y, m, d = today[0].split('-')

        yesterday = yl.split(',')
        r = daily_return(today[-1], yesterday[-1])

        returns.append((date(int(y), int(m), int(d)), r))

    return returns


def ex3():
    """
    Histogram the daily returns data

    If you can, fit a normal density to the data and plot that too
    """
    csv_list = read_data_file(DATA_PATH)
    returns_data = get_daily_returns(csv_list)
    hist_data = [t[1] for t in returns_data]

    # plot histogram
    pylab.xlabel('count')
    pylab.xlabel('daily % return')
    pylab.title('Nikkei Index Daily Returns Histogram')
    pylab.hist(hist_data, bins=40)
    pylab.show()

    # book example
    pylab.hist(hist_data, bins=200, normed=True)
    m, M = min(hist_data), max(hist_data)
    mu = pylab.mean(hist_data)
    sigma = pylab.std(hist_data)
    grid = pylab.linspace(m, M, 100)
    densityvalues = pylab.normpdf(grid, mu, sigma)
    pylab.plot(grid, densityvalues, 'r-')
    pylab.show()


# Exercise 4
def ex4():
    """
    Repeat Exercise 1, but using monthly data

        * Extract first quote of each month and plot as a time series
        * Note that first observation is not necessarily on the first day of month
              o first day of the month might be the weekend
    """
    csv_list = read_data_file(DATA_PATH)

    # set starting point
    first_line = csv_list[0]
    field = first_line.split(',')
    y, m, d = field[0].split('-')
    current_month = int(m)

    # collect month start quotes
    month_start_quotes = []
    for line in csv_list[1:]:
        field = line.split(',')
        quote_month = int(field[0].split('-')[1])
        if quote_month == current_month:
            continue

        # new month
        current_month = quote_month
        if line == csv_list[-1]:
            break
        month_start_quotes.append(field[-1])

    pylab.plot(month_start_quotes)
    pylab.xlabel('(1984-2009)')
    pylab.ylabel('monthly closing value')
    pylab.title('Nikkei Index Monthly Closings')
    pylab.show()


#
# MAIN
#
ex1()
ex2()
ex3()
ex4()
print '%s: ok' % (__file__)

